"""
通过python版opencv实现对数据增强的操作
"""
# ①　导入相关的包（os，cv2，torch等）
import numpy as np
import os
import cv2 as cv

# ②　创建一个文件夹，里面放置10张花的图片
# ③　将文件夹里面的图片分为训练集和测试集7:3
# ④　读取训练集内所有图片
print('Read pics')
imgs = []
names = []
dir = 'data/ori'
for file in os.listdir(dir):
    path = os.path.join(dir, file)
    names.append(file)
    img = cv.imread(path, cv.IMREAD_COLOR)
    imgs.append(img)
print('Reading over')

# ⑤　完成resize到224功能
# ⑥　结果正确
print('resize')
SIZE = 224
dest_dir = os.path.join(dir, '..', 'resize')
os.makedirs(dest_dir, exist_ok=True)
processed_imgs = []
for i, file in enumerate(names):
    img = imgs[i]
    H, W = img.shape[:2]
    if H > W:
        padding = H - W
        img = np.concatenate([img, np.zeros([H, padding, 3], dtype=np.uint8)], axis=1)
    elif W > H:
        padding = W - H
        img = np.concatenate([img, np.zeros([padding, W, 3], dtype=np.uint8)], axis=0)
    img = cv.resize(img, (SIZE, SIZE), interpolation=cv.INTER_CUBIC)
    processed_imgs.append(img)
    cv.imwrite(os.path.join(dest_dir, file), img)
print('resize over')

# ⑦　完成旋转30度
# ⑧　结果正确
print('Rotate')
dest_dir = os.path.join(dir, '..', 'rotate')
os.makedirs(dest_dir, exist_ok=True)
imgs = processed_imgs
processed_imgs = []
for i, file in enumerate(names):
    img = imgs[i]
    H, W = img.shape[:2]
    M = cv.getRotationMatrix2D((H // 2, W //2), 30, 1.)
    img = cv.warpAffine(img, M, (H, W))
    processed_imgs.append(img)
    cv.imwrite(os.path.join(dest_dir, file), img)
print('Rotate over')

# ⑨　完成线性对比度增强
# ⑩　结果正确

# 11　完成加入椒盐噪声数据增强
# 12　完成加入高斯噪声
# 13　结果正确
print('pepper and salt')
dest_dir = os.path.join(dir, '..', 'pepper_salt')
os.makedirs(dest_dir, exist_ok=True)
imgs = processed_imgs
processed_imgs = []
rate = 0.05
for i, file in enumerate(names):
    img = imgs[i]
    H, W = img.shape[:2]
    M = np.random.choice([0, 1, 2], [H, W, 3], True, [1 - rate * 2, rate, rate])
    img[M == 1] = 0
    img[M == 2] = 255
    processed_imgs.append(img)
    cv.imwrite(os.path.join(dest_dir, file), img)
print('pepper and salt over')
